{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7f8db942",
   "metadata": {},
   "source": [
    "# Configuring OpenAI LLM with PandasAI\n",
    "\n",
    "This notebook demonstrates how to configure PandasAI with OpenAI's LLM. \n",
    "\n",
    "We'll cover:\n",
    "- How to install the pandasai_openai llm extension\n",
    "- How to configure OpenAI LLM to be used with PandasAI\n",
    "- How to chat with one or multiple dataframes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84dc06e6",
   "metadata": {},
   "source": [
    "\n",
    "## Install openai llm extension\n",
    "\n",
    "PandasAI supports various [large language models](/v3/llms) (LLMs) to power its natural language query capabilities. OpenAI's models are available through a dedicated extension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e18747c7",
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "%pip install pandasai-openai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e6961c9",
   "metadata": {},
   "source": [
    "## Load OpenAI LLM"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56a3bb2b",
   "metadata": {},
   "source": [
    "Once you have installed the pandasai-openai extension, import the OpenAI class, and pass your OpenAI  API token as an argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ef36c4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandasai as pai\n",
    "from pandasai_openai import OpenAI\n",
    "\n",
    "llm = OpenAI(api_token=\"your_api_token\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3edfab53",
   "metadata": {},
   "source": [
    "(optional) You can check information about the llm model as it follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9061348d",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Print LLM configuration to verify setup\n",
    "print(f\"LLM Type: {type(llm)}\")\n",
    "print(f\"LLM Configuration:\")\n",
    "print(f\"- Model: {llm.model}\")  # Changed from model_name to model\n",
    "print(f\"- Temperature: {llm.temperature}\")\n",
    "print(f\"- Max Tokens: {llm.max_tokens}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7077745",
   "metadata": {},
   "source": [
    "## Configure OpenAI LLM in PandasAI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af21ef68",
   "metadata": {},
   "source": [
    "You can now set your OpenAI LLM as the default llm for PandasAI by simply passing the llm instance to the config via the `set()` method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb56211a",
   "metadata": {},
   "outputs": [],
   "source": [
    "pai.config.set({\n",
    "    \"llm\": llm,\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "710b19bf",
   "metadata": {},
   "source": [
    "(optional) Check the LLM has been correctly configured"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ba77740",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = pai.config.get()\n",
    "print(f\"- LLM in config: {type(config.llm).__name__}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5332309c",
   "metadata": {},
   "source": [
    "## Chat with your data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c9eb5b1",
   "metadata": {},
   "source": [
    "Once you have configured the llm, you can chat with one or multiple dataframes using the chat() method. For the purpose of this example, we are using a small dataset of heart disease patients from Kaggle (https://www.kaggle.com/datasets/arezaei81/heartcsv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cafc78d",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pai.read_csv(\"./data/heart.csv\")\n",
    "\n",
    "response = df.chat(\"What is the correlation between age and cholesterol?\")"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
